import torch
import torch.nn as nn


class PositionalEncoding2D(nn.Module):
    """2D位置编码，适用于棋盘等2D网格数据"""

    def __init__(self, channels, height, width, temperature=10000):
        super().__init__()
        self.channels = channels
        self.height = height
        self.width = width
        self.temperature = temperature

        # 预计算位置编码
        self.register_buffer('pos_encoding', self._get_positional_encoding())

    def _get_positional_encoding(self):
        # 为行和列分别创建位置编码
        pos_h = torch.arange(self.height).unsqueeze(1).float()
        pos_w = torch.arange(self.width).unsqueeze(0).float()

        # 创建频率
        dim_t = torch.arange(self.channels // 4).float()
        dim_t = self.temperature ** (2 * dim_t / (self.channels // 4))

        # 计算sin/cos编码
        pos_h = pos_h / dim_t
        pos_w = pos_w / dim_t

        # 组合行列编码
        pos_encoding = torch.zeros(self.channels, self.height, self.width)
        pos_encoding[0::4, :, :] = torch.sin(pos_h).unsqueeze(-1).expand(-1, -1, self.width)
        pos_encoding[1::4, :, :] = torch.cos(pos_h).unsqueeze(-1).expand(-1, -1, self.width)
        pos_encoding[2::4, :, :] = torch.sin(pos_w).unsqueeze(0).expand(self.height, -1, -1)
        pos_encoding[3::4, :, :] = torch.cos(pos_w).unsqueeze(0).expand(self.height, -1, -1)

        return pos_encoding.unsqueeze(0)  # 添加batch维度

    def forward(self, x):
        # x: [batch_size, channels, height, width]
        return x + self.pos_encoding[:, :x.size(1), :x.size(2), :x.size(3)]


class LearnablePositionalEncoding2D(nn.Module):
    """可学习的2D位置编码"""

    def __init__(self, channels, height, width):
        super().__init__()
        self.pos_embedding = nn.Parameter(
            torch.randn(1, channels, height, width) * 0.02
        )

    def forward(self, x):
        return x + self.pos_embedding
